Skip to content

fix(telemetry): drop use_span from _TracedStream.__anext__; eager-prime SSE under captured ctx#277

Merged
Kamilbenkirane merged 3 commits into
mainfrom
feat/fix-stream-context-detach
May 6, 2026
Merged

fix(telemetry): drop use_span from _TracedStream.__anext__; eager-prime SSE under captured ctx#277
Kamilbenkirane merged 3 commits into
mainfrom
feat/fix-stream-context-detach

Conversation

@Kamilbenkirane

Copy link
Copy Markdown
Member

Summary

  • _TracedStream.__anext__ wrapped await self._inner.__anext__() in with use_span(self._span, end_on_exit=False): so HTTPX child spans would parent to the GenAI span. Under async iteration that resumes across asyncio tasks, the attach/detach pair fired across mismatched contexts and raised ValueError: Token was created in a different Context. Closes _TracedStream.__anext__'s use_span context manager raises ValueError on detach across asyncio context boundaries #276.
  • Drop the with use_span block from __anext__ entirely. Capture an OTel context snapshot with the GenAI span active inside _stream (sync with use_span + contextvars.copy_context() — no await between attach and detach, no Task switch) and schedule the SSE iterator's first __anext__ as asyncio.create_task(coro, context=ctx). The HTTP request fires under that snapshot, so HTTPX auto-instrumentation captures the GenAI span as parent at request emission. Subsequent reads run in the consumer's own context and pull bytes from the open response — no new HTTPX spans, no per-chunk context manipulation.
  • Cancel the prime task if the consumer is cancelled before the first chunk arrives so the underlying open connection isn't leaked.
  • Single-file source change (client.py) plus the __anext__ cleanup (telemetry.py). Zero changes to provider mixins, http.py, _make_stream_request signature, Stream, or any public API.

Pattern aligns with opentelemetry-instrumentation-openai-v2's BaseStreamWrapper.__anext__, which has zero use_span / attach / detach calls.

References #272 (where the workaround originated).

Test plan

  • make lint clean
  • make format --check clean
  • make typecheck clean
  • make test — 637 unit tests pass (5 new tests in tests/unit_tests/test_telemetry_streaming.py covering: __anext__ does not invoke use_span, prime preserves event order, empty stream yields nothing, first-pull error propagates, first pull runs under captured ctx with subsequent pulls under caller's ctx, cancel-mid-first-pull cancels the prime task)
  • Real-call smoke under a multi-chunk streaming consumer: confirm zero "Failed to detach context" warnings, confirm HTTPX child spans still parent to the GenAI span, confirm TTFC / usage / response-model / finish-reason / metric histograms still emit

…me SSE under captured ctx

`_TracedStream.__anext__` wrapped `await self._inner.__anext__()` in
`with use_span(self._span, end_on_exit=False):` to make the GenAI span
the parent of HTTPX child spans. Under FastAPI StreamingResponse the
attach/detach pair fired across asyncio tasks, raising
`ValueError: Token was created in a different Context` ~19x per chat
session. Closes #276.

Capture an OTel context snapshot with the GenAI span active in `_stream`
(sync `with use_span`, no await between attach and detach). Schedule the
SSE iterator's first __anext__ as `asyncio.create_task(coro, context=ctx)`
so the HTTP request fires under the snapshot; HTTPX auto-instrumentation
captures the GenAI span as parent at request emission. Subsequent reads
run in the consumer's own context and pull bytes from the open response.

Drop the `with use_span` block from `_TracedStream.__anext__`.
…irst chunk

Without this, a cancellation interrupting `await task` in
`_prime_with_context` left the underlying create_task running in the
background, holding the open SSE/HTTP connection. Add a `BaseException`
catch that calls `task.cancel()` then re-raises. Add regression test
covering cancel-mid-first-pull.
@claude

claude Bot commented May 6, 2026

Copy link
Copy Markdown

Code review

No issues found. Checked for bugs and CLAUDE.md compliance.

…py to telemetry.py

The first pass parked an async-task + contextvars helper at the top of
client.py and added asyncio + contextvars imports there. That leaked
telemetry mechanism into the module that defines the SDK base layering.

Move the helper to telemetry.py (where it joins trace_stream / record_*
/ use_span and where contextvars belongs); rename to
bind_first_pull_to_span; revert client.py imports and drop the inline
ctx-capture block from `_stream`. `_stream` regains its single-verb
orchestration rhythm with one new line:

    sse_iterator = telemetry.bind_first_pull_to_span(sse_iterator, span)

Same behavior, same tests (signatures updated to pass a span instead of
a Context; new test asserts the first pull sees the span as current
OTel context and subsequent pulls don't).
@Kamilbenkirane Kamilbenkirane merged commit 32b05cc into main May 6, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

_TracedStream.__anext__'s use_span context manager raises ValueError on detach across asyncio context boundaries

1 participant